#include "Tracker.h"

short doubleToShort(double d)
{
	if (d>32766) return 32766;
	if (d<-32767) return -32767;
	return (short(d));
}

void Tracker::setFirstFrame(const Mat frame, Rect box)
{
	cv::Mat matMask = cv::Mat::zeros(frame.size(), CV_8UC1);
	rectangle(matMask, box, Scalar(255), -1);
	detector->detectAndCompute(frame, matMask, first_kp, first_desc);
	object_bb.push_back(cv::Point2f(static_cast<float>(box.x), static_cast<float>(box.y)));
	object_bb.push_back(cv::Point2f(static_cast<float>(box.x + box.width), static_cast<float>(box.y)));
	object_bb.push_back(cv::Point2f(static_cast<float>(box.x + box.width), static_cast<float>(box.y + box.height)));
	object_bb.push_back(cv::Point2f(static_cast<float>(box.x), static_cast<float>(box.y + box.height)));
}
Mat Tracker::process(const Mat frame)
{
	vector<KeyPoint> kp;
	Mat desc;
	detector->detectAndCompute(frame, noArray(), kp, desc);
	vector<vector<DMatch>> matches;
	vector<KeyPoint> matched1, matched2;
	matcher->knnMatch(first_desc, desc, matches, 2);
	for (unsigned i = 0; i < matches.size(); i++)
	{
		if (matches[1].empty())
			continue;
		if (matches[i][0].distance < nn_match_ratio * matches[i][1].distance)
		{
			matched1.push_back(first_kp[matches[i][0].queryIdx]);
			matched2.push_back(kp[matches[i][0].trainIdx]);
		}
	}
	Mat inlier_mask, homography;
	vector<KeyPoint> inliers1, inliers2;
	vector<DMatch> inlier_matches;
	if (matched1.size() >= 4)
	{
		homography = findHomography(Points(matched1), Points(matched2),
			RANSAC, ransac_thresh, inlier_mask);
	}
	if (matched1.size() < 4 || homography.empty())
		return Mat();
	for (unsigned i = 0; i < matched1.size(); i++)
	{
		if (inlier_mask.at<uchar>(i))
		{
			int new_i = static_cast<int>(inliers1.size());
			inliers1.push_back(matched1[i]);
			inliers2.push_back(matched2[i]);
			inlier_matches.push_back(DMatch(new_i, new_i, 0));
		}
	}
	vector<Point2f> new_bb;
	perspectiveTransform(object_bb, new_bb, homography);
	if (inliers1.size() >= bb_min_inliers)
	{
		drawBoundingBox(frame, new_bb);
	}
	else return Mat();
	return homography;
}

void Tracker::drawBoundingBox(Mat image, vector<Point2f> bb)
{
	for (unsigned i = 0; i < bb.size()-1; i++) {
		line(image, bb[i], bb[i + 1], Scalar(0, 0, 255), 2);
	}
	line(image, bb[bb.size()-1], bb[0], Scalar(0, 0, 255), 2);
}

vector<Point2f> Tracker::Points(vector<KeyPoint> keypoints)
{
	vector<Point2f> res;
	for (unsigned i = 0; i < keypoints.size(); i++) {
		res.push_back(keypoints[i].pt);
	}
	return res;
}
